from funasr import AutoModel
import re
import time
import threading
from config import punc_max_cache_size, punc_use_cache

# 全局模型实例
global_punc_model = None
# 结果缓存 - 用于存储已处理过的文本结果，避免重复处理相同的文本
result_cache = {}
# 缓存大小限制，防止内存占用过大
MAX_CACHE_SIZE = punc_max_cache_size
# 线程锁，用于保护缓存访问
cache_lock = threading.Lock()

def load_punc_model():
    """
    加载标点恢复模型
    
    Returns:
        加载成功返回模型实例，失败返回None
    """
    # cache_dir或cache或punc_model均不能控制；默认下载到：C:\Users\EDY\.cache\modelscope\hub\models\iic\punc_ct-transformer_cn-en-common-vocab471067-large\
    try:
        print("加载标点恢复模型...")
        model = AutoModel(model="ct-punc", disable_update=True)
        print("标点恢复模型加载完成")
        return model
    except Exception as e:
        print(f"标点恢复模型加载失败: {str(e)}")
        return None

def simplify_punctuation(text):
    """将重复的标点符号简化为单个标点符号"""
    # 处理中文标点
    text = re.sub(r'[？]+', '？', text)  # 多个问号变一个
    text = re.sub(r'[！]+', '！', text)  # 多个感叹号变一个
    text = re.sub(r'[。]+', '。', text)  # 多个句号变一个
    text = re.sub(r'[，]+', '，', text)  # 多个逗号变一个
    text = re.sub(r'[；]+', '；', text)  # 多个分号变一个
    text = re.sub(r'[：]+', '：', text)  # 多个冒号变一个
    
    # 处理英文标点
    text = re.sub(r'[?]+', '?', text)   # 多个英文问号变一个
    text = re.sub(r'[!]+', '!', text)   # 多个英文感叹号变一个
    text = re.sub(r'[.]+', '.', text)   # 多个英文句号变一个
    text = re.sub(r'[,]+', ',', text)   # 多个英文逗号变一个
    text = re.sub(r'[;]+', ';', text)   # 多个英文分号变一个
    text = re.sub(r'[:]+', ':', text)   # 多个英文冒号变一个
    
    return text

def clear_result_cache():
    """清空结果缓存"""
    global result_cache
    with cache_lock:
        result_cache.clear()
    print("标点恢复结果缓存已清空")

def contains_chinese(text, min_ratio=0.3):
    """中文字符占比判断"""
    chinese_chars = re.findall(r'[\u4e00-\u9fa5]', text)
    total_chars = len(text)
    if total_chars == 0:
        return False
    chinese_ratio = len(chinese_chars) / total_chars
    return chinese_ratio >= min_ratio  # 中文占比≥min_ratio%才返回True

def chinese_marks_punctuation_replace(text, min_ratio=0.3):
    """功能：对文本末尾的标点进行智能替换（仅当中文占比大于min_ratio时生效）。
    替换规则：
    单标点 → 转换为对应的中文标点（如 . → 。、? → ？）。
    重复标点 → 按优先级选择最合适的标点（！ ？ 。 > ； > ， 、 ：）。
    非中文文本（中文占比小于min_ratio） → 不进行任何替换。
    """
    # 去除首尾空白字符
    text = text.strip()
    
    # 如果中文占比<60%，直接返回原文本
    if not contains_chinese(text):
        return text
    
    # 英文标点 -> 中文标点映射
    en_to_zh_punct = {
        '.': '。',
        '?': '？',
        '!': '！',
        ';': '；',
        ',': '，',
        ':': '：'
    }
    
    # 优先级组（从高到低）
    priority_groups = [
        ['。', '？', '！'],  # 最高优先级
        ['；'],             # 中等优先级
        ['，', '、', '：']   # 最低优先级
    ]
    
    # 匹配末尾的标点（中文或英文，至少1个）
    match = re.search(r'([.?!;,:。？！；，、：]+)$', text)
    if not match:
        return text  # 无标点，直接返回
    
    original_puncts = match.group(1)
    
    # 1. 先转换英文标点为中文标点
    converted_puncts = []
    for p in original_puncts:
        converted_puncts.append(en_to_zh_punct.get(p, p))  # 如果是英文标点则转换，否则保留
    
    # 2. 如果只有一个标点，直接替换
    if len(converted_puncts) == 1:
        new_text = text[:match.start()] + converted_puncts[0]
        return new_text
    
    # 3. 如果是重复标点，按优先级选择最合适的
    replacement = None
    for group in priority_groups:
        for p in converted_puncts:
            if p in group:
                replacement = p
                break
        if replacement is not None:
            break
    
    # 如果没有匹配任何优先级组（理论上不会发生）
    if replacement is None:
        replacement = converted_puncts[-1]
    
    # 替换末尾标点
    new_text = text[:match.start()] + replacement
    return new_text


def text_punc_restore(text, model=None, use_cache=True):
    """
    对文本进行标点恢复，并处理重复标点
    
    Args:
        text: 输入文本
        model: 可选，标点恢复模型实例
        use_cache: 是否使用缓存，默认为True
        
    Returns:
        处理后的文本
    """
    if not text or text.strip() == "":
        return text
    
    # 如果启用缓存且文本在缓存中存在，直接返回缓存结果
    if use_cache:
        with cache_lock:
            if text in result_cache:
                return result_cache[text]
    
    # 如果没有传入模型实例，则使用全局模型
    global global_punc_model
    if model is None:
        model = global_punc_model
    
    # 确保模型不为None
    if model is None:
        print("警告: 标点恢复模型未提供且全局模型未加载，返回原始文本")
        return text
    
    # 执行标点恢复
    res = model.generate(input=text)
    res_text = res[0]['text']
    
    
    # 如果启用缓存，将结果存入缓存
    if use_cache:
        with cache_lock:
            # 如果缓存太大，清除一部分
            if len(result_cache) >= MAX_CACHE_SIZE:
                # 简单策略：清除一半的缓存
                keys_to_remove = list(result_cache.keys())[:MAX_CACHE_SIZE // 2]
                for key in keys_to_remove:
                    result_cache.pop(key, None)
            
            # 将当前结果加入缓存
            result_cache[text] = res_text
    
    return res_text

# FastAPI 相关代码，仅在作为主模块运行且需要API功能时使用
# 为了避免在导入模块时创建路由器，将这些代码提取到一个函数中
def create_punc_router():
    # 创建标点恢复路由
    from fastapi import APIRouter, Body
    from pydantic import BaseModel

    # 创建标点恢复路由器
    punc_router = APIRouter()

    class TextRequest(BaseModel):
        text: str

    @punc_router.post("/restore")
    async def restore_punctuation(request: TextRequest):
        """
        将文本进行标点恢复处理
        
        - 为无标点文本添加适当的标点
        - 自动处理重复的标点符号（如 ??? 将变为 ?）
        - 返回原始文本和处理后的文本
        """
        try:
            result = text_punc_restore(request.text, model=global_punc_model)
            return {"original": request.text, "result": result}
        except Exception as e:
            print(f"标点恢复API处理失败: {str(e)}")
            # 返回原始文本，避免API失败
            return {"original": request.text, "result": request.text, "error": str(e)}
    
    return punc_router

# 创建供导出的路由器实例
punc_router = create_punc_router()
    
if __name__ == "__main__":
    # 测试函数
    # 预加载模型，避免多次加载
    print("加载标点恢复模型用于测试...")
    
    # 记录模型加载时间
    load_start_time = time.time()
    test_model = load_punc_model()
    load_end_time = time.time()
    print(f"模型加载耗时: {load_end_time - load_start_time:.4f} 秒")
    
    test_cases = [
        "支持多种语言，包括中文英语日语韩语法语德语西班牙语阿拉伯语等74种语言",
        "徐霞客去过哪些地方？"
    ]
    
    # 总处理时间统计
    total_process_time = 0
    total_chars = 0
    total_cases = len(test_cases)
    
    print("\n===== 开始性能测试 =====\n")
    
    # 首先进行一次预热调用，避免首次调用的时间不准确
    print("预热调用...")
    _ = text_punc_restore("预热测试文本", model=test_model)
    print("预热完成")
    print("-" * 50)
    
    for i, test in enumerate(test_cases, 1):
        chars = len(test)
        total_chars += chars
        
        print(f"测试 [{i}/{total_cases}] (长度: {chars} 字符)")
        print(f"原文: {test}")
        
        # 记录单次处理开始时间
        start_time = time.time()
        result = text_punc_restore(test, model=test_model)
        end_time = time.time()
        
        # 计算处理时间
        process_time = end_time - start_time
        total_process_time += process_time
        
        print(f"结果: {result}")
        print(f"处理用时: {process_time:.4f} 秒")
        print(f"每字符耗时: {(process_time / chars):.6f} 秒/字符")
        print("-" * 50)
    
    # 计算平均处理时间
    if total_cases > 0 and total_chars > 0:
        avg_time = total_process_time / total_cases
        avg_time_per_char = total_process_time / total_chars
        
        print("\n===== 性能测试结果摘要 =====")
        print(f"总处理文本数: {total_cases} 条")
        print(f"总字符数: {total_chars} 字符")
        print(f"总处理用时: {total_process_time:.4f} 秒")
        print(f"平均每条文本处理用时: {avg_time:.4f} 秒")
        print(f"平均每字符处理用时: {avg_time_per_char:.6f} 秒/字符")
        
        # 估算每分钟处理字符数
        chars_per_minute = (total_chars / total_process_time) * 60
        print(f"估算处理速度: {chars_per_minute:.2f} 字符/分钟")
        
        print("\n===== 测试完成 =====")

